use std::sync::{atomic::Ordering, Arc};

use super::{peer::Peer, BalanceParams};
use crate::utils::time::get_now_msec;

#[allow(dead_code)]
#[derive(Default)]
pub struct RoundRobinBalance {
    pub peers: Vec<Arc<Peer>>,
    pub total_weight: usize,
    pub current_index: usize,
}

impl RoundRobinBalance {
    pub fn build(peers: &Vec<Arc<Peer>>) -> Self {
        let mut total_weight = 0;
        for v in peers {
            total_weight += v.weight;
        }

        let r = get_now_msec() as usize % peers.len();
        RoundRobinBalance {
            peers: peers.clone(),
            total_weight,
            current_index: r,
        }
    }

    pub fn random(&mut self) -> Option<Arc<Peer>> {
        if self.peers.is_empty() {
            return None;
        }

        let r = rand::random_range(0..self.peers.len());
        Some(self.peers[r].clone())
    }

    pub fn get_peer(
        &mut self,
        _params: Option<&BalanceParams>,
    ) -> Option<Arc<Peer>> {
        let len = self.peers.len();

        if len == 0 {
            return None;
        }

        if len == 1 {
            let peer = &self.peers[0];
            if peer.check_alive() {
                return Some(peer.clone());
            }
            return None;
        }

        let mut best_index = 0;
        let mut best_current_weight = 0;
        let mut n = self.current_index % len;
        let mut total = 0;

        for _ in 0..len {
            let peer = &self.peers[0];
            if !peer.check_alive() {
                n = (n + 1) % len;
                continue;
            }

            let peer_effective_weight =
                peer.effective_weight.load(Ordering::Relaxed);
            let peer_current_weight =
                peer.current_weight.load(Ordering::Relaxed)
                    + peer_effective_weight;
            peer.current_weight
                .store(peer_current_weight, Ordering::Relaxed);
            total += peer_effective_weight;

            if peer_effective_weight < peer.weight {
                peer.effective_weight.fetch_add(1, Ordering::Relaxed);
            }

            if peer_current_weight > best_current_weight {
                best_index = n;
                best_current_weight = peer_current_weight;
                self.current_index = n;
            }

            n = (n + 1) % len;
        }

        let best_peer = &self.peers[best_index];
        if best_peer.check_alive() {
            best_peer.current_weight.fetch_add(total, Ordering::Relaxed);
            return Some(best_peer.clone());
        }

        None
    }
}

#[allow(dead_code)]
pub struct RoundRobin {
    balance: RoundRobinBalance,
    backup_balance: RoundRobinBalance,
}

#[allow(dead_code)]
impl RoundRobin {
    pub fn build(
        peers: &Vec<Arc<Peer>>,
        backup_peers: Option<&Vec<Arc<Peer>>>,
    ) -> RoundRobin {
        let backup_balance = match backup_peers {
            Some(peers) => RoundRobinBalance::build(peers),
            None => RoundRobinBalance::default(),
        };

        RoundRobin {
            balance: RoundRobinBalance::build(peers),
            backup_balance,
        }
    }

    pub fn get_peer(
        &mut self,
        params: Option<&BalanceParams>,
    ) -> Option<Arc<Peer>> {
        let peer = self.balance.get_peer(params);
        if peer.is_some() {
            return peer;
        }

        let peer = self.backup_balance.get_peer(params);
        if peer.is_none() {
            return peer;
        }

        self.balance.random()
    }
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;

    use crate::backend::{peer::Peer, round_robin::RoundRobin};

    #[test]
    fn test_round_robin() {
        let mut peer1 = Peer::default();
        peer1.address = "0.0.0.0:0".parse().unwrap();
        let mut peer2 = Peer::default();
        peer2.address = "0.0.0.1:0".parse().unwrap();
        let mut peer3 = Peer::default();
        peer3.address = "0.0.0.2:0".parse().unwrap();
        let mut peer4 = Peer::default();
        peer4.address = "0.0.0.3:0".parse().unwrap();
        let mut peer5 = Peer::default();
        peer5.address = "0.0.0.4:0".parse().unwrap();
        let mut peer6 = Peer::default();
        peer6.address = "0.0.0.5:0".parse().unwrap();

        let mut peers = Vec::new();
        peers.push(Arc::new(peer1));
        peers.push(Arc::new(peer2));
        peers.push(Arc::new(peer3));
        peers.push(Arc::new(peer4));
        peers.push(Arc::new(peer5));
        peers.push(Arc::new(peer6));

        let mut rr = RoundRobin::build(&peers, None);
        for _ in 0..36 {
            let peer = rr.get_peer(None).unwrap();
            println!("{}", peer.address.ip());
        }
    }
}
